import pinocchio as pin
import numpy as np
from numpy.linalg import norm, inv, pinv


urdf_filename = "piper_description/urdf/piper_description.urdf"  # 替换为你的URDF文件路径

model = pin.buildModelFromUrdf(urdf_filename)
data = model.createData()

def print_all_frames():
    
    # 打印所有frames的信息
    print(f"模型中共有 {len(model.frames)} 个frames:")
    print("\nID\t名称\t\t类型\t\t父关节ID\t相对位置")
    print("-" * 80)

    for frame_id, frame in enumerate(model.frames):
        # Frame类型的描述
        frame_type = {
            pin.FrameType.OP_FRAME: "操作框架",
            pin.FrameType.JOINT: "关节",
            pin.FrameType.BODY: "链接",
            pin.FrameType.FIXED_JOINT: "固定关节"
        }.get(frame.type, "未知类型")
        
        # 父关节ID
        parent_joint_id = frame.parent
        
        # 相对位置（相对于父关节的位置）
        placement = frame.placement
        position = placement.translation
        rotation = placement.rotation
        
        # 打印信息（限制名称长度为15个字符以保持格式整齐）
        name = frame.name[:15] + ("..." if len(frame.name) > 15 else "")
        print(f"{frame_id}\t{name:<18}{frame_type:<15}{parent_joint_id:<15}"
            f"({position[0]:.3f}, {position[1]:.3f}, {position[2]:.3f})")


# 选择末端执行器的frame ID（根据你的机器人模型修改）
end_effector_name = "ee"  # 替换为你的末端执行器名称
end_effector_id = model.getFrameId(end_effector_name)



# 设置关节配置（例如：全零配置） go zero!
q = np.zeros(model.nq)

# 执行前向运动学计算
pin.forwardKinematics(model, data, q)
pin.updateFramePlacements(model, data)  # 更新所有frame的位置


# 通过frame ID获取特定frame的全局位置
frame_id = model.getFrameId("link8")  # arm?
frame_pose = data.oMf[frame_id]  # SE3变换矩阵
frame_position = frame_pose.translation  # 位置向量 [x, y, z]

rotation_matrix = frame_pose.rotation
euler_angles = pin.rpy.matrixToRpy(rotation_matrix)
euler_degrees = np.rad2deg(euler_angles)

print(f"Frame '{model.frames[frame_id].name}' 的全局位置: {frame_position}, rxyz={euler_degrees}")


## Try to use ik

target_position = p = np.array([ 408.264,   70.954,   16.008])
# target_rotation = np.eye(3)
target_rotation = pin.rpy.rpyToMatrix(np.array([ 3.124,  0.175, -2.97]))

target_pose = pin.SE3(target_rotation, target_position)

# 初始关节配置（可以是随机的或当前配置）
q_init = np.zeros(model.nq)  # 使用零配置作为初始值

# 逆运动学求解参数
max_iterations = 10000  # 最大迭代次数
tolerance = 1e-6       # 收敛容差
dt = 0.01               # 步长
damping = 10       # 阻尼因子

# 逆运动学求解主循环
q = q_init.copy()
for i in range(max_iterations):
    # 前向运动学计算当前末端执行器位姿
    pin.forwardKinematics(model, data, q)
    pin.updateFramePlacements(model, data)
    current_pose = data.oMf[end_effector_id]
    
    # 计算位姿误差
    # 位置误差
    position_error = target_position - current_pose.translation
    
    # 姿态误差（使用对数映射）
    rotation_error = pin.log3(current_pose.rotation.T @ target_rotation)
    
    # 组合误差向量
    error = np.concatenate([position_error, rotation_error])
    
    # 检查收敛性
    if norm(error) < tolerance:
        print(f"Converged after {i} iterations!")
        break
    
    # 计算雅可比矩阵
    # 注意：必须先调用computeJointJacobians和updateFramePlacements
    pin.computeJointJacobians(model, data, q)
    jacobian = pin.getFrameJacobian(
        model, data, end_effector_id, pin.LOCAL_WORLD_ALIGNED
    )
    
    # 使用阻尼最小二乘法求解关节速度
    JTJ = jacobian.T @ jacobian
    damping_matrix = damping * np.eye(model.nv)
    joint_velocity = pinv(jacobian.T @ jacobian + damping_matrix) @ jacobian.T @ error
    
    # 更新关节配置
    q = pin.integrate(model, q, dt * joint_velocity)
    
    # 可选：添加关节限制检查和处理
    
    # 检查是否收敛
pin.forwardKinematics(model, data, q)
pin.updateFramePlacements(model, data)
final_pose = data.oMf[end_effector_id]
final_position = final_pose.translation
final_rotation = final_pose.rotation

print("目标位置:", target_position)
print("最终位置:", final_position)
print("位置误差:", norm(target_position - final_position))

print("\n目标姿态矩阵:")
print(target_rotation)
print("\n最终姿态矩阵:")
print(final_rotation)

# 计算最终的位姿误差
final_position_error = target_position - final_position
final_rotation_error = pin.log3(final_rotation.T @ target_rotation)
final_error = np.concatenate([final_position_error, final_rotation_error])
print("\n最终总误差:", norm(final_error))

# 打印最终的关节配置
print("\n逆运动学解（关节角度）:", q)

joints = np.rad2deg(q[:6])

joints = q[:6]

###############################################################







import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D


# 目标位置（x, y, z）
target_position = p = np.array([ 408.264,   70.954,   16.008])
# target_rotation = np.eye(3)
# 目标姿态（使用欧拉角表示，然后转换为旋转矩阵）
target_euler_angles = np.array([ 3.124,  0.175, -2.97])  # [rx, ry, rz] (弧度)
target_rotation = pin.rpy.rpyToMatrix(target_euler_angles)

# 组合成SE3变换矩阵
target_pose = pin.SE3(target_rotation, target_position)

# 初始关节配置（可以使用零配置或当前配置）
q_init = np.zeros(model.nq)

# IK求解参数
max_iterations = 1000
tolerance = 1e-6
dt = 0.1
damping = 10

# 存储迭代过程（用于可视化）
q_history = [q_init.copy()]
error_history = []

# 获取关节限制
lower_limits = np.array([model.lowerPositionLimit[i] for i in range(model.nq)])
upper_limits = np.array([model.upperPositionLimit[i] for i in range(model.nq)])

# 逆运动学求解主循环
q = q_init.copy()
for i in range(max_iterations):
    # 前向运动学计算当前末端执行器位姿
    pin.forwardKinematics(model, data, q)
    pin.updateFramePlacements(model, data)
    current_pose = data.oMf[end_effector_id]
    
    # 计算位姿误差
    position_error = target_position - current_pose.translation
    rotation_error = pin.log3(current_pose.rotation.T @ target_rotation)
    error = np.concatenate([position_error, rotation_error])
    
    # 记录误差
    error_history.append(norm(error))
    
    # 检查收敛性
    if norm(error) < tolerance:
        print(f"收敛于迭代次数 {i}")
        break
    
    # 计算雅可比矩阵
    pin.computeJointJacobians(model, data, q)
    jacobian = pin.getFrameJacobian(
        model, data, end_effector_id, pin.LOCAL_WORLD_ALIGNED
    )
    
    # 使用阻尼最小二乘法求解关节速度
    joint_velocity = pinv(jacobian.T @ jacobian + damping * np.eye(model.nv)) @ jacobian.T @ error
    
    # 更新关节配置
    q = pin.integrate(model, q, dt * joint_velocity)
    # q = np.clip(q, lower_limits, upper_limits)  # 截断到限制范围内
    q_history.append(q.copy())
    
else:
    print(f"达到最大迭代次数，最终误差: {norm(error)}")

# 计算最终结果
pin.forwardKinematics(model, data, q)
pin.updateFramePlacements(model, data)
final_pose = data.oMf[end_effector_id]
final_position = final_pose.translation
final_rotation = final_pose.rotation
final_euler = pin.rpy.matrixToRpy(final_rotation)

# 打印结果
print("\n==== 逆运动学求解结果 ====")
print(f"目标位置: {target_position}")
print(f"最终位置: {final_position}")
print(f"位置误差: {norm(target_position - final_position)}")

print(f"\n目标欧拉角 [rx, ry, rz]: {np.degrees(target_euler_angles)}°")
print(f"最终欧拉角 [rx, ry, rz]: {np.degrees(final_euler)}°")
print(f"姿态误差: {norm(pin.log3(final_rotation.T @ target_rotation))}")

print(f"\n逆运动学解（关节角度）: {q}")

# 可视化误差收敛过程
plt.figure(figsize=(10, 6))
plt.plot(error_history)
plt.title('IK求解误差收敛过程')
plt.xlabel('迭代次数')
plt.ylabel('误差范数')
plt.yscale('log')
plt.grid(True)
plt.show()

# 可视化3D位置轨迹（可选）
if len(q_history) > 1:
    positions = []
    for q_step in q_history:
        pin.forwardKinematics(model, data, q_step)
        pin.updateFramePlacements(model, data)
        positions.append(data.oMf[end_effector_id].translation)
    
    positions = np.array(positions)
    
    fig = plt.figure(figsize=(10, 8))
    ax = fig.add_subplot(111, projection='3d')
    
    # 绘制轨迹
    ax.plot(positions[:, 0], positions[:, 1], positions[:, 2], 'b-', label='IK轨迹')
    
    # 标记起点和终点
    ax.scatter(positions[0, 0], positions[0, 1], positions[0, 2], c='r', marker='o', s=100, label='起点')
    ax.scatter(positions[-1, 0], positions[-1, 1], positions[-1, 2], c='g', marker='o', s=100, label='终点')
    
    # 标记目标位置
    ax.scatter(target_position[0], target_position[1], target_position[2], c='m', marker='x', s=100, label='目标')
    
    ax.set_xlabel('X')
    ax.set_ylabel('Y')
    ax.set_zlabel('Z')
    ax.set_title('末端执行器位置轨迹')
    ax.legend()
    plt.show()







# import mujoco_py
# import os
# mj_path = mujoco_py.utils.discover_mujoco()
# xml_path = os.path.join(mj_path, 'model', 'piper_ros/src/piper_description/mujoco_model/piper_description.xml')
# model = mujoco_py.load_model_from_path(xml_path)
# sim = mujoco_py.MjSim(model)
# print(sim.data.qpos)
# sim.step()
# print(sim.data.qpos)